Cpu memory graph break #3886
Cpu memory graph break #3886
Conversation
github-actions
bot
added
component: tests
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
There are some changes that do not conform to Python style guidelines:
--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/_compiler.py 2025-11-04 20:05:23.825034+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/_compiler.py 2025-11-04 20:05:55.253944+00:00
@@ -876,15 +876,14 @@
# This is done to release CPU memory.
for attr in dir(gm):
if attr.startswith("_frozen_param"):
delattr(gm, attr)
-
-
from torch_tensorrt.dynamo.conversion._ConverterRegistry import DYNAMO_CONVERTERS
+
DYNAMO_CONVERTERS.disallowed_targets = set()
-
+
for name, _ in partitioned_module.named_children():
submodule = getattr(partitioned_module, name)
# filter on the GraphModule
if not isinstance(submodule, torch.fx.graph_module.GraphModule):
continue
narendasan
left a comment
narendasan
left a comment
There was a problem hiding this comment.
Do you have a test case or something to demonstrate this feature?
|
|
||
| logger = logging.getLogger(__name__) | ||
| NON_BREAKABLE_OP_LISTS = [ | ||
| ["addmm", "addmm"], |
There was a problem hiding this comment.
Just a note for implementation later.
- this should use an actual subgraph definition
- it should use pytorch op targets not strings
addmmshould be decomposed right so the graph we want ismm->add- There should be a user facing API to modify this list similar to what we have for passes
|
|
||
| def calculate_num_of_break(self, subgraphs: List[Subgraph]) -> int: | ||
|
|
||
| def calculate_size_budget( |
There was a problem hiding this comment.
Should there be an API to define this manually?
There was a problem hiding this comment.
Yeah I think so. For now you can just hardcode and play with it
cehongwang
force-pushed
the
cpu-memory-graph-break
branch
from
7f0e504 to
18ccadf
Compare
|
|
Improve usability by automating nn.Module -> atomic fx graph |
cehongwang
force-pushed
the
cpu-memory-graph-break
branch
from
18ccadf to
f03ab2c
Compare
| return x | ||
|
|
||
|
|
||
| All_FUSION_PATTERNS = [ |
There was a problem hiding this comment.
We could cache the graphs if we do symbolic trace on register
There was a problem hiding this comment.
Do you mean to trace the graphs when the program starts every time? Do you think that would cause unnecessary latency when cpu memory is enough. I am thinking maybe we could use LRU cache or something so it will only be called once and it's lazy initialization
| L2_LIMIT_FOR_TILING = -1 | ||
| USE_DISTRIBUTED_MODE_TRACE = False | ||
| OFFLOAD_MODULE_TO_CPU = False | ||
| CPU_MEMORY_BUDGET = -1 |
There was a problem hiding this comment.
Use an optional instead since this is not a TRT api we dont need -1 to mean let us decide
| return psutil.Process().memory_info().rss / 1024 / 1024 | ||
|
|
||
|
|
||
| def release_memory() -> None: |
| torch._dynamo.reset() | ||
|
|
||
|
|
||
| def compile_one(idx: int, ir: str): |
|
|
||
| def size_of_subgraphs(self, subgraphs: List[Subgraph]) -> List[int]: | ||
| """ | ||
| This function calculates the size of the subgraph. |
There was a problem hiding this comment.
Can you describe the algorithms here so we have reference for later?
cehongwang
force-pushed
the
cpu-memory-graph-break
branch
from
2389d34 to
7f9373f
Compare
cehongwang
force-pushed
the
cpu-memory-graph-break
branch
from
7f9373f to
9ee7e67
Compare
Description
Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.
Fixes # (issue)
Type of change
Please delete options that are not relevant and/or add your own.
Checklist: